"""
股票相关的Celery任务
"""
import logging
import datetime
from celery import group, chain
from celery.utils.log import get_task_logger
from sqlalchemy import func

from modules.tasks.celery_app import app
from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from modules.data_collector.stock_collector import StockCollector
from modules.data_processor.indicators_calculator import IndicatorsCalculator

# 配置任务日志
logger = get_task_logger(__name__)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def update_stock_list(self):
    """更新股票列表信息"""
    logger.info("开始更新股票列表信息")
    try:
        collector = StockCollector()
        new_count, update_count = collector.collect_stock_info()
        logger.info(f"股票列表更新完成: 新增 {new_count} 个，更新 {update_count} 个")
        return {
            "status": "success",
            "new_count": new_count,
            "update_count": update_count
        }
    except Exception as e:
        logger.error(f"更新股票列表失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_stock_data(self, security_code, start_date=None, end_date=None):
    """
    采集单只股票的历史数据
    
    Args:
        security_code: 股票代码
        start_date: 开始日期, 格式为YYYYMMDD
        end_date: 结束日期, 格式为YYYYMMDD
    """
    logger.info(f"开始采集股票 {security_code} 的历史数据")
    
    try:
        # 创建任务记录
        db = next(get_db())
        task_record = CollectionTask(
            task_id=self.request.id,
            task_type='stock_daily',
            security_code=security_code,
            start_date=datetime.datetime.strptime(start_date, '%Y%m%d').date() if start_date else None,
            end_date=datetime.datetime.strptime(end_date, '%Y%m%d').date() if end_date else None,
            status='pending'
        )
        db.add(task_record)
        db.commit()
        db.close()
        
        # 执行数据采集
        collector = StockCollector()
        count = collector.collect_daily_quote(security_code, start_date, end_date, self.request.id)
        
        logger.info(f"股票 {security_code} 采集完成，获取 {count} 条数据")
        
        # 如果有新数据，触发技术指标计算
        if count > 0:
            # 使用链式任务计算技术指标
            calculate_stock_indicators.delay(security_code)
        
        return {
            "status": "success",
            "security_code": security_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {security_code} 数据失败: {str(e)}")
        # 更新任务记录状态
        try:
            db = next(get_db())
            task = db.query(CollectionTask).filter(CollectionTask.task_id == self.request.id).first()
            if task:
                task.status = 'failed'
                task.error_message = str(e)
                task.update_time = datetime.datetime.now()
                db.commit()
        except Exception as db_e:
            logger.error(f"更新任务状态失败: {str(db_e)}")
        finally:
            db.close()
            
        raise self.retry(exc=e)


@app.task(bind=True, queue='indicators', max_retries=2, retry_backoff=30)
def calculate_stock_indicators(self, security_code, recalculate=False):
    """
    计算单只股票的技术指标
    
    Args:
        security_code: 股票代码
        recalculate: 是否重新计算所有指标
    """
    logger.info(f"开始计算股票 {security_code} 的技术指标")
    
    try:
        calculator = IndicatorsCalculator()
        count = calculator.calculate_security_indicators(security_code, recalculate)
        
        logger.info(f"股票 {security_code} 技术指标计算完成，处理了 {count} 条数据")
        return {
            "status": "success",
            "security_code": security_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"计算股票 {security_code} 技术指标失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data')
def schedule_daily_stock_update(self):
    """
    调度每日股票数据更新任务
    每天收盘后执行，获取所有活跃股票的最新数据
    """
    logger.info("开始调度每日股票数据更新任务")
    
    try:
        # 首先更新股票列表
        update_stock_list.delay()
        
        # 获取当前日期和前一交易日
        today = datetime.datetime.now().strftime('%Y%m%d')
        yesterday = (datetime.datetime.now() - datetime.timedelta(days=1)).strftime('%Y%m%d')
        
        # 获取所有活跃股票
        db = next(get_db())
        stocks = db.query(SecurityInfo).filter(
            SecurityInfo.type == 'stock',
            SecurityInfo.is_active == 1
        ).all()
        db.close()
        
        logger.info(f"准备更新 {len(stocks)} 只股票的日线数据")
        
        # 创建并行任务组
        tasks = []
        for stock in stocks:
            tasks.append(collect_stock_data.s(stock.code, yesterday, today))
        
        # 分批执行任务，避免一次创建过多任务
        batch_size = 50
        for i in range(0, len(tasks), batch_size):
            batch_tasks = tasks[i:i+batch_size]
            group(batch_tasks).apply_async()
            logger.info(f"已调度 {i+len(batch_tasks)}/{len(tasks)} 个股票数据采集任务")
        
        return {
            "status": "success",
            "message": f"已调度 {len(stocks)} 只股票的数据更新任务"
        }
    except Exception as e:
        logger.error(f"调度股票更新任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='stock_data')
def collect_stock_historical_data(self, security_code, years=3):
    """
    采集股票的历史数据
    
    Args:
        security_code: 股票代码
        years: 采集最近几年的数据，默认3年
    """
    logger.info(f"开始采集股票 {security_code} 的 {years} 年历史数据")
    
    try:
        # 计算起始日期
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=365*years)).strftime('%Y%m%d')
        
        # 执行数据采集
        result = collect_stock_data.delay(security_code, start_date, end_date)
        
        return {
            "status": "success",
            "message": f"已调度股票 {security_code} 的历史数据采集任务",
            "task_id": result.id
        }
    except Exception as e:
        logger.error(f"调度股票历史数据采集任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='indicators')
def recalculate_all_indicators(self, security_type='stock'):
    """
    重新计算所有证券的技术指标
    
    Args:
        security_type: 证券类型，默认为'stock'
    """
    logger.info(f"开始重新计算所有{security_type}的技术指标")
    
    try:
        calculator = IndicatorsCalculator()
        processed_count = calculator.calculate_all_securities_indicators(
            security_type=security_type,
            recalculate=True
        )
        
        logger.info(f"技术指标重新计算完成，处理了 {processed_count} 个证券")
        return {
            "status": "success",
            "security_type": security_type,
            "processed_count": processed_count
        }
    except Exception as e:
        logger.error(f"重新计算技术指标失败: {str(e)}")
        raise


@app.task(bind=True, queue='stock_data')
def check_data_integrity(self, security_type='stock', days=30):
    """
    检查数据完整性，找出最近一段时间缺失数据的证券
    
    Args:
        security_type: 证券类型，默认为'stock'
        days: 检查最近几天的数据，默认30天
    """
    logger.info(f"开始检查最近 {days} 天的{security_type}数据完整性")
    
    try:
        db = next(get_db())
        
        # 获取所有活跃证券
        securities = db.query(SecurityInfo).filter(
            SecurityInfo.type == security_type,
            SecurityInfo.is_active == 1
        ).all()
        
        # 计算日期范围
        end_date = datetime.datetime.now().date()
        start_date = end_date - datetime.timedelta(days=days)
        
        # 检查数据完整性
        incomplete_securities = []
        
        for security in securities:
            # 获取该证券在日期范围内的记录数
            count = db.query(func.count(DailyQuote.id)).filter(
                DailyQuote.security_id == security.id,
                DailyQuote.trade_date >= start_date,
                DailyQuote.trade_date <= end_date
            ).scalar()
            
            # 计算工作日数量（简化处理，按照75%工作日计算）
            expected_days = int(days * 0.75)
            
            # 如果记录数少于预期，则认为数据不完整
            if count < expected_days:
                incomplete_securities.append({
                    "code": security.code,
                    "name": security.name,
                    "expected": expected_days,
                    "actual": count
                })
        
        logger.info(f"数据完整性检查完成，发现 {len(incomplete_securities)} 个证券数据不完整")
        
        # 对于数据不完整的证券，调度任务重新采集
        for security in incomplete_securities:
            collect_stock_historical_data.delay(security["code"], 1)  # 采集最近1年数据
        
        db.close()
        
        return {
            "status": "success",
            "security_type": security_type,
            "incomplete_count": len(incomplete_securities),
            "incomplete_securities": incomplete_securities
        }
    except Exception as e:
        logger.error(f"数据完整性检查失败: {str(e)}")
        raise


# 测试函数
def test_stock_tasks():
    """测试股票任务函数"""
    # 更新股票列表
    update_result = update_stock_list.delay()
    print(f"更新股票列表任务已提交，任务ID: {update_result.id}")
    
    # 等待任务完成
    result = update_result.get(timeout=120)
    print(f"更新股票列表结果: {result}")
    
    # 采集单只股票数据
    db = next(get_db())
    stock = db.query(SecurityInfo).filter(
        SecurityInfo.type == 'stock'
    ).first()
    db.close()
    
    if stock:
        # 计算最近30天的日期范围
        end_date = datetime.datetime.now().strftime('%Y%m%d')
        start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
        
        collect_result = collect_stock_data.delay(stock.code, start_date, end_date)
        print(f"股票 {stock.code} 数据采集任务已提交，任务ID: {collect_result.id}")
        
        # 等待任务完成
        result = collect_result.get(timeout=120)
        print(f"股票数据采集结果: {result}")
    else:
        print("没有找到股票数据，无法测试采集任务")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_stock_tasks()